#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@ author:hengk
@ contact: hengk@foxmail.com
@ datetime:2019-10-25 15:22
"""

#!/usr/env/bin python3

"""
Generate training and test images.
"""
import os

# prevent opencv use all cpus
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'

import cv2
import traceback
import numpy as np
import multiprocessing as mp
from itertools import repeat
from tenacity import retry

from libs.utils import load_chars
from libs.timer import Timer
from libs.config import load_config
from parse_args import parse_args
from barcode_renderer.renderer import Renderer








def start_listen(q, fname):
    """ listens for messages on the q, writes to file. """

    f = open(fname, mode='a', encoding='utf-8')
    while 1:
        m = q.get()
        if m == STOP_TOKEN:
            break
        try:
            f.write(str(m) + '\n')
        except:
            traceback.print_exc()

        with lock:
            if counter.value % 1000 == 0:
                f.flush()
    f.close()


@retry
def gen_img_retry(renderer, img_index):
    try:
        return renderer.gen_img(img_index)
    except Exception as e:
        print("Retry gen_img: %s" % str(e))
        traceback.print_exc()
        raise Exception


def generate_img(img_index, q=None):
    # global flags, lock, counter
    # Make sure different process has different random seed
    np.random.seed()
    image,code = gen_img_retry(renderer, img_index)

    base_name = '{:08d}'.format(img_index)
    fname = os.path.join(flags.save_dir, base_name + '.jpg')

    cv2.imwrite(fname, image)

    label = "{} {}".format(base_name+".jpg", code)

    if q is not None:
        q.put(label)

    with lock:
        counter.value += 1
        print_end = '\n' if counter.value == flags.num_img else '\r'
        if counter.value % 100 == 0 or counter.value == flags.num_img:
            print("{}/{} {:2d}%".format(counter.value,
                                        flags.num_img,
                                        int(counter.value / flags.num_img * 100)),
                  end=print_end)




def restore_exist_labels(label_path):
    # 如果目标目录存在 labels.txt 则向该目录中追加图片
    start_index = 0
    if os.path.exists(label_path):
        start_index = len(load_chars(label_path))
        print('Generate more text images in %s. Start index %d' % (flags.save_dir, start_index))
    else:
        print('Generate text images in %s' % flags.save_dir)
    return start_index


def get_num_processes(flags):
    processes = flags.num_processes
    if processes is None:
        processes = max(os.cpu_count(), 2)
    return processes


if __name__ == "__main__":

    lock = mp.Lock()
    counter = mp.Value('i', 0)
    STOP_TOKEN = 'kill'

    flags = parse_args()

    renderer = Renderer(flags)
    #设置图片目录，标签目录
    label_path = os.path.join(flags.output_dir, 'labels.txt')
    flags.save_dir = os.path.join(flags.output_dir, flags.images_dir)
    if not os.path.exists(flags.save_dir):
        os.makedirs(flags.save_dir)

    manager = mp.Manager()
    q = manager.Queue()

    start_index = restore_exist_labels(label_path)

    timer = Timer(Timer.SECOND)
    timer.start()
    with mp.Pool(processes=get_num_processes(flags)) as pool:

        pool.apply_async(start_listen, (q, label_path))
        pool.starmap(generate_img, zip(range(start_index, start_index + flags.num_img), repeat(q)))

        q.put(STOP_TOKEN)
        pool.close()
        pool.join()
    timer.end("Finish generate data")


